#include <gtest/gtest.h>
#include "../../kvcache-service/include/KVCacheService.h"
#include "../../kvcache-memorypool/include/MemoryPool.h"
#include "proto/kvcache.grpc.pb.h"
#include "Common.h"
#include "KvcBaseDefine.h"
#include <grpcpp/grpcpp.h>

using grpc::ServerContext;
using namespace kvcache;

class KVCacheServiceTest : public ::testing::Test {
protected:
  void SetUp() override {
  }

  void TearDown() override {
  }

  
};

TEST_F(KVCacheServiceTest, PrefetchPrefixCacheSuccess) {
    KVCacheServiceImpl service;
    service.init();
    PrefetchPrefixCacheRequest request;
    auto entry = request.add_entries();
    for (int64_t i = 1; i <= 64; i++) {
        uint64_t localHashId = (uint64_t)i;
        auto blockHashId = entry->add_blockhashids();
        blockHashId->set_hashid(i);
        blockHashId->set_parenthashid(i);
        blockHashId->set_localhashid(localHashId);
    }
    request.set_tpnum(1);

    auto cache_info = request.mutable_cacheinfo();
    cache_info->set_startlayer(0);
    cache_info->set_endlayer(1);
    cache_info->set_layernum(1);
    cache_info->set_layercachesize(33554432);

    request.set_memsize(33554432);
    PrefetchPrefixCacheResponse response;
    auto start = getCurrentTimestampUs();
    grpc::Status status = service.PrefetchPrefixCache(nullptr, &request, &response);
    printf("[cost time] %lu ms\n", (getCurrentTimestampUs() - start) / MICROSECONDS_PER_MILLISECOND);
    EXPECT_EQ(response.result(), 0);
}

/* 如需运行此用例，设定memorypool大小为160MB
 * cmake .. -DBUILD_TEST=1 -DDEFAULT_MEMPOOL_SIZE=167772160 && make -j
 * ./test/kvcache-service/ServiceTest --gtest_filter=*TestTwoPromptConflict_NoEvictableBlkCauseAllocFail*
 */
TEST_F(KVCacheServiceTest, TestTwoPromptConflict_NoEvictableBlkCauseAllocFail) {
    int prompt_len = 4;
    int resp_len = 4;
    int mem_size = 33554432;

    KVCacheServiceImpl service;
    service.init();
    MemoryPool *pool = service._pools[mem_size].get();

    PrefetchPrefixCacheRequest prefetchReq[2];
    for (int i = 0; i < 2; i++) {
        auto entry = prefetchReq[i].add_entries();
        for (int j = 0; j < prompt_len; j++) {
            auto blockHashId = entry->add_blockhashids();
            int64_t hashId = i * 65536 + j;
            uint64_t localHashId = (uint64_t)hashId;
            blockHashId->set_hashid(hashId);
            blockHashId->set_parenthashid(hashId);
            blockHashId->set_localhashid(localHashId);
        }
        prefetchReq[i].set_tpnum(1);

        auto cache_info = prefetchReq[i].mutable_cacheinfo();
        cache_info->set_startlayer(0);
        cache_info->set_endlayer(64);
        cache_info->set_layernum(64);
        cache_info->set_layercachesize(262144);

        prefetchReq[i].set_memsize(mem_size);
    }

    QueryPrefixCacheRequest queryReq[2];
    for (int i = 0; i < 2; i++) {
        for (int j = 0; j < prompt_len; j++) {
            queryReq[i].add_hashids(i * 65536 + j);
        }
    }
    
    AllocateMemoryRequest allocReq[2];
    for (int i = 0; i < 2; i++) {
        for (int j = 0; j < prompt_len; j++) {
            auto blockHashIdAttr = allocReq[i].add_blockhashidattrs();
            int64_t hashId = i * 65536 + 16384 + j;
            uint64_t localHashId = (uint64_t)hashId;
            blockHashIdAttr->set_hashid(hashId);
            blockHashIdAttr->set_parenthashid(hashId);
            blockHashIdAttr->set_localhashid(localHashId);
        }
        allocReq[i].set_memsize(mem_size);
        allocReq[i].set_writeinrecordsize(1);

        allocReq[i].mutable_cacheinfo()->set_startlayer(0);
        allocReq[i].mutable_cacheinfo()->set_endlayer(64);
        allocReq[i].mutable_cacheinfo()->set_layernum(64);
        allocReq[i].mutable_cacheinfo()->set_layercachesize(262144);
    }

    sealMemoryRequest sealReq[2];
    for (int i = 0; i < 2; i++) {
        for (int j = 0; j < prompt_len; j++) {
            sealReq[i].add_hashids(i * 65536 + 16384 + j);
        }
        sealReq[i].set_dosave(true);
    }

    /* 模拟2个会话并发，第2个会话导致第1个会话中部分segment被逐出的场景，过程如下：
     * 1）会话1 PrefetchPrefixCache -> QueryPrefixCache -> allocateMemory -> sealMemory
     * 2）会话2 PrefetchPrefixCache -> QueryPrefixCache -> allocateMemory -> sealMemory
     * 3) 打印lru状态，预期会话1中部分segment被逐出
     */
    PrefetchPrefixCacheResponse prefetchRsp;
    QueryPrefixCacheResponse queryRsp;
    AllocateMemoryResponse allocRsp;
    sealMemoryResponse sealRsp;
    grpc::Status status;
    for (int i = 0; i < 2; i++) {
        status = service.PrefetchPrefixCache(nullptr, &prefetchReq[i], &prefetchRsp);
        std::cout << "[PrefetchPrefixCache] i: " << i << ", status: " << status.ok() << std::endl;
        if (prefetchRsp.result() != KVC_OK) {
            std::cout << "[PrefetchPrefixCache] BREAK! ret: " << prefetchRsp.result() << std::endl;
            break;
        }

        status = service.QueryPrefixCache(nullptr, &queryReq[i], &queryRsp);
        std::cout << "[QueryPrefixCache] i: " << i << ", status: " << status.ok() << std::endl;

        status = service.AllocateMemory(nullptr, &allocReq[i], &allocRsp);
        std::cout << "[AllocateMemory] i: " << i << ", status: " << status.ok() << std::endl;

        status = service.sealMemory(nullptr, &sealReq[i], &sealRsp);
        std::cout << "[sealMemory] i: " << i << ", status: " << status.ok() << std::endl;
        if (sealRsp.result() != KVC_OK) {
            std::cout << "[sealMemory] BREAK! ret: " << sealRsp.result() << std::endl;
            break;
        }

        Segment *segment;
        for (int32_t j = 0; j < prompt_len; j++) {
            pool->getSegmentByID(i * 65536 + 16384 + j, &segment);
            segment->setIsSaved(true);
        }
    }

    std::vector<BlockHashId> hashIds = pool->getAllCachedSegmentIds();
    printf("All CachedSegmentId: ");
    for (auto &id : hashIds) {
        printf("%ld ", id.hashId);
    }
    printf("\n");
    // 开大页情况下
    ASSERT_EQ(hashIds.size(), 0);
    // 不开大页情况下
    // ASSERT_EQ(hashIds.size(), 4);
    // ASSERT_EQ(hashIds[0].hashId, 16384);
    // ASSERT_EQ(hashIds[1].hashId, 16385);
    // ASSERT_EQ(hashIds[2].hashId, 16386);
    // ASSERT_EQ(hashIds[3].hashId, 16387);
}

/* 如需运行此用例，设定memorypool大小为160MB
 * cmake .. -DBUILD_TEST=1 -DDEFAULT_MEMPOOL_SIZE=167772160 && make -j
 * ./test/kvcache-service/ServiceTest --gtest_filter=*TestTwoPromptConflict_EnoughEvictableBlk*
 */
TEST_F(KVCacheServiceTest, TestTwoPromptConflict_EnoughEvictableBlk) {
    int prompt_len = 4;
    int resp_len = 4;
    int mem_size = 33554432;

    KVCacheServiceImpl service;
    service.init();
    MemoryPool *pool = service._pools[mem_size].get();

    PrefetchPrefixCacheRequest prefetchReq[2];
    for (int i = 0; i < 2; i++) {
        auto entry = prefetchReq[i].add_entries();
        for (int j = 0; j < prompt_len; j++) {
            auto blockHashId = entry->add_blockhashids();
            int64_t hashId = i * 65536 + j;
            uint64_t localHashId = (uint64_t)hashId;
            blockHashId->set_hashid(hashId);
            blockHashId->set_parenthashid(hashId);
            blockHashId->set_localhashid(localHashId);
        }
        prefetchReq[i].set_tpnum(1);

        auto cache_info = prefetchReq[i].mutable_cacheinfo();
        cache_info->set_startlayer(0);
        cache_info->set_endlayer(64);
        cache_info->set_layernum(64);
        cache_info->set_layercachesize(262144);

        prefetchReq[i].set_memsize(mem_size);
    }

    QueryPrefixCacheRequest queryReq[2];
    for (int i = 0; i < 2; i++) {
        for (int j = 0; j < prompt_len; j++) {
            queryReq[i].add_hashids(i * 65536 + j);
        }
    }
    
    AllocateMemoryRequest allocReq[2];
    for (int i = 0; i < 2; i++) {
        for (int j = 0; j < prompt_len; j++) {
            auto blockHashIdAttr = allocReq[i].add_blockhashidattrs();
            int64_t hashId = i * 65536 + 16384 + j;
            uint64_t localHashId = (uint64_t)hashId;
            blockHashIdAttr->set_hashid(hashId);
            blockHashIdAttr->set_parenthashid(hashId);
            blockHashIdAttr->set_localhashid(localHashId);
        }
        allocReq[i].set_memsize(mem_size);
        allocReq[i].set_writeinrecordsize(1);

        allocReq[i].mutable_cacheinfo()->set_startlayer(0);
        allocReq[i].mutable_cacheinfo()->set_endlayer(64);
        allocReq[i].mutable_cacheinfo()->set_layernum(64);
        allocReq[i].mutable_cacheinfo()->set_layercachesize(262144);
    }

    sealMemoryRequest sealReq[2];
    for (int i = 0; i < 2; i++) {
        for (int j = 0; j < prompt_len; j++) {
            sealReq[i].add_hashids(i * 65536 + 16384 + j);
        }
    }

    /* 模拟2个会话并发，第2个会话导致第1个会话中部分segment被逐出的场景，过程如下：
     * 1）会话1 PrefetchPrefixCache -> QueryPrefixCache -> allocateMemory -> sealMemory
     * 2) 等待2s
     * 3）会话2 PrefetchPrefixCache -> QueryPrefixCache -> allocateMemory -> sealMemory
     * 4) 打印lru状态，预期会话1中部分segment被逐出
     */
    PrefetchPrefixCacheResponse prefetchRsp;
    QueryPrefixCacheResponse queryRsp;
    AllocateMemoryResponse allocRsp;
    sealMemoryResponse sealRsp;
    grpc::Status status;
    for (int i = 0; i < 2; i++) {
        status = service.PrefetchPrefixCache(nullptr, &prefetchReq[i], &prefetchRsp);
        std::cout << "[PrefetchPrefixCache] i: " << i << ", status: " << status.ok() << std::endl;
        if (prefetchRsp.result() != KVC_OK) {
            std::cout << "[PrefetchPrefixCache] BREAK! ret: " << prefetchRsp.result() << std::endl;
            break;
        }

        status = service.QueryPrefixCache(nullptr, &queryReq[i], &queryRsp);
        std::cout << "[QueryPrefixCache] i: " << i << ", status: " << status.ok() << std::endl;

        status = service.AllocateMemory(nullptr, &allocReq[i], &allocRsp);
        std::cout << "[AllocateMemory] i: " << i << ", status: " << status.ok() << std::endl;

        status = service.sealMemory(nullptr, &sealReq[i], &sealRsp);
        std::cout << "[sealMemory] i: " << i << ", status: " << status.ok() << std::endl;
        if (sealRsp.result() != KVC_OK) {
            std::cout << "[sealMemory] BREAK! ret: " << sealRsp.result() << std::endl;
            break;
        }

        Segment *segment;
        for (int32_t j = 0; j < prompt_len; j++) {
            pool->getSegmentByID(i * 65536 + 16384 + j, &segment);
            segment->setIsSaved(true);
        }

        sleep(2);
    }

    std::vector<BlockHashId> hashIds = pool->getAllCachedSegmentIds();
    printf("All CachedSegmentId: ");
    for (auto &id : hashIds) {
        printf("%ld ", id.hashId);
    }
    printf("\n");
    // 开大页情况下
    ASSERT_EQ(hashIds.size(), 0);
    // 不开大页情况下
    // ASSERT_EQ(hashIds.size(), 6);
    // ASSERT_EQ(hashIds[0].hashId, 81920);
    // ASSERT_EQ(hashIds[1].hashId, 81921);
    // ASSERT_EQ(hashIds[2].hashId, 81922);
    // ASSERT_EQ(hashIds[3].hashId, 81923);
    // ASSERT_EQ(hashIds[4].hashId, 16384);
    // ASSERT_EQ(hashIds[5].hashId, 16385);
}

/* 如需运行此用例，设定memorypool大小为160MB
 * cmake .. -DBUILD_TEST=1 -DDEFAULT_MEMPOOL_SIZE=167772160 && make -j
 * ./test/kvcache-service/ServiceTest --gtest_filter=*TestReadFromStorageFail*
 */
TEST_F(KVCacheServiceTest, TestReadFromStorageFail) {
    KVCacheServiceImpl service;
    service.init();

    int prompt_len = 8;
    int resp_len = 4;
    int mem_size = 33554432;
    PrefetchPrefixCacheRequest prefetchReq;
    auto entry = prefetchReq.add_entries();
    for (int j = 0; j < prompt_len; j++) {
        auto blockHashId = entry->add_blockhashids();
        uint64_t localHashId = (uint64_t)j;
        blockHashId->set_hashid(j);
        blockHashId->set_parenthashid(j);
        blockHashId->set_localhashid(localHashId);
    }
    prefetchReq.set_tpnum(1);

    auto cache_info = prefetchReq.mutable_cacheinfo();
    cache_info->set_startlayer(0);
    cache_info->set_endlayer(64);
    cache_info->set_layernum(64);
    cache_info->set_layercachesize(262144);

    prefetchReq.set_memsize(mem_size);

    /* 模拟1个会话并发，下盘读失败触发releaseSegment, freelist长度应该为6，即全部空闲可用*/
    PrefetchPrefixCacheResponse prefetchRsp;
    grpc::Status status = service.PrefetchPrefixCache(nullptr, &prefetchReq, &prefetchRsp);
    std::cout << "[PrefetchPrefixCache] status: " << status.ok() << std::endl;
    assert(prefetchRsp.result() == KVC_OK);

    MemoryPool *pool = service._pools[mem_size].get();
    std::vector<BlockHashId> hashIds = pool->getAllCachedSegmentIds();
    assert(hashIds.empty());
    // 开大页情况下
    ASSERT_EQ(pool->getFreelist().size(), 0);
    // 不开大页情况下
    // ASSERT_EQ(pool->getFreelist().size(), 6);
}

int main(int argc, char** argv) {
    testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}